import sys

import boto3
from pyathena import connect

out_loc = 's3://lambda-athena-log'
aki = 'AKIAWTBG3E42IMYODWHL'
sak = 'onuA0VOzcCvidVPaZGcVESHrLSz2hZ9Qr5gGPG0c'
rn = 'cn-northwest-1'

glue = boto3.client('glue', region_name=rn, aws_access_key_id=aki, aws_secret_access_key=sak)
athena = boto3.client('athena', region_name=rn, aws_access_key_id=aki, aws_secret_access_key=sak)


def run(source, target):
    print(source, '->', target)
    sp = source.split('.')
    source_db = sp[0]
    source_table = sp[1]

    tp = target.split('.')
    target_db = tp[0]
    target_table = tp[1]

    table_desc = glue.get_table(
        DatabaseName=source_db,
        Name=source_table,
    )
    partition_keys = []
    for part in table_desc['Table']['PartitionKeys']:
        partition_keys.append(part['Name'])

    paginator = glue.get_paginator('get_partitions')
    cursor = connect(s3_staging_dir=out_loc,
                     region_name=rn,
                     aws_access_key_id=aki,
                     aws_secret_access_key=sak
                     ).cursor(schema_name=target_db)

    page_iterator = paginator.paginate(
        DatabaseName=source_db,
        TableName=source_table,
        ExcludeColumnSchema=True)

    i = 1
    sql = f'ALTER TABLE {target} ADD '
    for response in page_iterator:
        partitions = response['Partitions']
        for partition in partitions:
            i += 1
            values = partition['Values']
            filters = []
            for j in range(0, len(partition_keys)):
                filters.append(f'''{partition_keys[j]} = '{values[j]}' ''')
            filter_sql = ','.join(filters)
            sql = f'{sql} PARTITION ({filter_sql}) '
            if i % 99 == 0:
                print('添加', i, '个分区')
                sql = sql + ';'
                try:
                    print(f"执行{sql}")
                    cursor.execute(sql)
                    cursor.fetchall()
                    sql = f'ALTER TABLE {target} ADD '
                except Exception as e:
                    print(sql, '错误')
                    print(e)
    try:
        print(f"执行{sql}")
        cursor.execute(sql)
        cursor.fetchall()
    except Exception as e:
        print(sql, '错误')
        print(e)
    print('共添加', i, '个分区')


if __name__ == '__main__':
    event = sys.argv
    source_table = sys.argv[1]
    target_table = sys.argv[2]
    run(source_table, target_table)
    print("任务完成")
